import os

from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_community.chat_models import ChatTongyi
from langsmith import traceable

# 从环境变量中获取阿里云百练的 API Key
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
# 阿里云百练的官网地址
DASHSCOPE_API_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"


class State(TypedDict):
    messages: Annotated[list, add_messages]


chat_model = ChatTongyi(model="qwen-max")


# 定义聊天机器人的节点函数，接收当前状态并返回更新的消息列表
def chat_bot(state: State):
    return {"messages": [chat_model.invoke(state["messages"])]}


@traceable(project_name="langgraph_chat_bot_app")
def run_chat_bot():
    graph_builder = StateGraph(State)
    graph_builder.add_node("chat_bot", chat_bot)
    graph_builder.add_edge(START, "chat_bot")
    graph_builder.add_edge("chat_bot", END)
    graph = graph_builder.compile()
    while True:
        # 获取用户输入
        user_input = input("User: ")

        # 可以随时通过输入 "quit"、"exit" 或 "q" 退出聊天循环
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        for event in graph.stream({"messages": ("user", user_input)}):
            # 遍历每个事件的值
            for value in event.values():
                # 打印输出 chat_bot 生成的最新消息
                print("Assistant:", value["messages"][-1].content)


if __name__ == '__main__':
    run_chat_bot()
